import os
import sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import json
import base64

from configs.config import load_config


def process_labelme_json(old_json_path, img_path, save_path):
    # 读取旧的 json
    with open(old_json_path, "r", encoding="utf-8") as f:
        data = json.load(f)
    # 读取图片并编码为 base64
    with open(img_path, "rb") as img_file:
        img_base64 = base64.b64encode(img_file.read()).decode("utf-8")
    # 构建新 json
    new_data = {
        "version": "5.2.1",
        "flags": {},
        "shapes": []
    }
    
    for i, shape in enumerate(data["shapes"]):
        if i % 1 == 0: # !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
            new_shape = {
                "label": shape["label"],
                "points": shape["points"]
            }
            new_data["shapes"].append(new_shape)
    # 复制其它字段
    for key in ["imagePath", "imageHeight", "imageWidth"]:
        if key in data:
            new_data[key] = data[key]
    # 填充 imageData
    new_data["imageData"] = img_base64
    # 保存新 json
    with open(save_path, "w", encoding="utf-8") as f:
        json.dump(new_data, f, ensure_ascii=False, indent=2)

def main(label_json_root, image_root):
    # 遍历所有形如 data331_b58 的子文件夹
    for folder in os.listdir(label_json_root):
        if folder == '0':
            continue
        print(folder)
        label_folder = os.path.join(label_json_root, folder)
        if os.path.isdir(label_folder):
            # 寻找 _mosaic.json 文件
            json_filename = f"{folder}_mosaic.json"
            old_json_path = os.path.join(label_folder, json_filename)
            if not os.path.exists(old_json_path):
                print(f"跳过：{old_json_path} 不存在")
                continue
            # 去图片目录找图片
            img_folder = os.path.join(image_root, folder)
            img_filename = f"total_BEV_{folder}_all.png"
            img_path = os.path.join(img_folder, img_filename)
            if not os.path.exists(img_path):
                print(f"跳过：{img_path} 不存在")
                continue
            # 保存到图片目录，命名 data331_b58.json
            save_path = os.path.join(img_folder, f"{folder}.json")
            process_labelme_json(old_json_path, img_path, save_path)
        break

if __name__ == "__main__":
    cfg = load_config("./configs/data_process.py")

    label_json_root = r"D:\desktop\data"
    image_root = r"./workspace/result/bev"
    main(label_json_root, image_root)
